import os

import cv2
import numpy as np
from PIL import Image
from tqdm import tqdm
from pathlib import Path
import matplotlib.pyplot as plt


# 指定单张图像路径
# img_dir = Path(r'img_dir\train')
# mask_dir = Path(r'ann_dir\train')
img_path = Path(r'/home/luoluoluo/data/dataset/elevator/rope_voc/JPEGImages/058_10605.jpg')
mask_path = Path(r'/home/luoluoluo/data/dataset/elevator/rope_voc/SegmentationClass/058_10605.png')
# for img_path, mask_path in zip(img_dir.glob('*.jpg'), mask_dir.glob('*.png')):
img = cv2.imread(str(img_path))
mask = cv2.imread(str(mask_path))
print(np.unique(mask))
print(f"开始读写{img_path}")

palette = [
    ['background', [0,0,0]],
    ['pet', [0,255,0]],
    ['rope', [0,0,255]],
    ['human', [144,238,144]],
]

palette_dict = {}
for idx, each in enumerate(palette):
    palette_dict[idx] = each[1]

mask = mask[:,:,0]

# 将整数ID，映射为对应类别的颜色
viz_mask_bgr = np.zeros((mask.shape[0], mask.shape[1], 3))
for idx in palette_dict.keys():
    viz_mask_bgr[np.where(mask==idx)] = palette_dict[idx]
viz_mask_bgr = viz_mask_bgr.astype('uint8')

# 将语义分割标注图和原图叠加显示
opacity = 0.2 # 透明度越大，可视化效果越接近原图
label_viz = cv2.addWeighted(img, opacity, viz_mask_bgr, 1-opacity, 0)

plt.figure(figsize=(10, 6))
plt.imshow(label_viz[:,:,::-1])
plt.axis('off')
plt.show()
